#include "Mapping.h"

//---------------MATCH HYPOTHESIS-------------//

/**
*  This is the constructor called when the MH is first made. There will always be two
*  lines mapped, corresponding to the two lines in the relationship matrices (in a LineGroup).
*/
MH::MH(int belief, Line g_line_0, Line g_line_1, Line c_line_0, Line c_line_1){
    goalMappings = new int[N_GOAL_LINES];
    mappingEvidence = new int[N_GOAL_LINES];
    g_lines = new Line[N_GOAL_LINES];
    c_lines = new Line[N_GOAL_LINES];
    for(int i=0; i < N_GOAL_LINES; i++){
        if(i == g_line_0.getId()){
            goalMappings[i] = c_line_0.getId();
            mappingEvidence[i] = belief;
        }else if(i == g_line_1.getId()){
            goalMappings[i] = c_line_1.getId();
            mappingEvidence[i] = belief;
        }else{
            goalMappings[i] = -1;
            mappingEvidence[i] = 0;
        }
    }
    g_lines[0] = g_line_0;
    g_lines[1] = g_line_1;
    c_lines[0] = c_line_0;
    c_lines[1] = c_line_1;
    n_lines = 2;
    scoreMH();
}

/**
*  This constructor is called when MHs are merged. All of the merged arrays etc. are
*  worked out before this is called.
*/
MH::MH(int *mappingEvidence, int* goalMappings, Line *g_lines, Line *c_lines, int n_lines){
    this->mappingEvidence = mappingEvidence;
    this->goalMappings = goalMappings;
    this->g_lines = g_lines;
    this->c_lines = c_lines;
    this->n_lines = n_lines;    
    scoreMH();
}

int MH::getN_lines(){
    return n_lines;
}

int *MH::getMappings(){
    return goalMappings;
}

int MH::getMappingAt(int index){
    return goalMappings[index];
}

int *MH::getAllEvidence(){
    return mappingEvidence;
}

int MH::getEvidenceAt(int index){
    return mappingEvidence[index];
}

Line *MH::getAllG_lines(){
    return g_lines;
}

Line MH::getG_lineAt(int index){
    return g_lines[index];
}

Line *MH::getAllC_lines(){
    return c_lines;
}

Line MH::getC_lineAt(int index){
    return c_lines[index];
}

double MH::getScore(){
    return score;
}

/**
*  Scores the MH by comparing lines to each other (relative to the mean line attributes of all
*  lines mapped in the MH so far).
*/
void MH::scoreMH(){
    int belief_sum = 0;
    double score_tot = 0.0;
    //double *g_means = calculateMeans(g_lines, n_lines);    
    //double *c_means = calculateMeans(c_lines, n_lines);
    for(int i=0; i<n_lines; i++){
        int goal_id = g_lines[i].getId();
        int belief_i = mappingEvidence[goal_id];
        score_tot += lineSimScore(g_lines[i], c_lines[i], G_MEANS, C_MEANS)*belief_i;
        belief_sum += belief_i;
    }
    score = score_tot/belief_sum;
}

double *calculateMeans(Line *lines, int n_lines){
    int x_sum = 0;
    int y_sum = 0;
    double orient_sum = 0.0;
    double length_sum = 0.0;
    for(int i=0; i<n_lines; i++){
        Line l_i = lines[i];
        int *l_i_points = l_i.getPoints();
        int *l_i_mid = addrToCoords(l_i_points[2]);
        x_sum += l_i_mid[0];
        y_sum += l_i_mid[1];
        orient_sum += l_i.getOrientation();
        length_sum += l_i.getLength();
    }
    double *means = new double[3];
    means[0] = coordsToAddr((x_sum/n_lines), (y_sum/n_lines));
    means[1] = orient_sum/n_lines;
    means[2] = length_sum/n_lines;
    return means;
}

/**
*  Checks for overlap in the MH Line mappings, new information,
*  and lack of repetition to say whether two MHs can merge.
*/
bool canMerge(MH MH1, MH MH2){
    bool overlaps = 0;
    for(int i=0; i<N_GOAL_LINES; i++){
        int MH1_i = MH1.getMappingAt(i);
        int MH2_i = MH2.getMappingAt(i);
        if((MH1_i != -1) && (MH2_i != -1)) {
            if(MH1_i == MH2_i)
                overlaps = 1;
            else                //they are in conflict :(
                return 0;
        }
    }
    bool noRepeats = 1;
    for(int k=0; k<N_GOAL_LINES; k++){
        int MH1_k = MH1.getMappingAt(k);
        if(MH1_k == -1)
            continue;
        for(int m=0; m<N_GOAL_LINES; m++){
            int MH2_m =  MH2.getMappingAt(m);
            if((MH1_k == MH2_m) && (k!=m)){
                noRepeats = 0;
            }
        }
    }
    return (!(MUST_OVERLAP && !overlaps) && noRepeats);
}

/**
*  Takes two MHs and merges them into one.
*/
MH mergeMHs(MH MH_0, MH MH_1){
    int *mappings_merge = new int[N_GOAL_LINES];
    int *evidence_merge = new int[N_GOAL_LINES];
    Line *g_lines_merge = MH_0.getAllG_lines();
    Line *c_lines_merge = MH_0.getAllC_lines();
    int n_lines_merge = 0;
    for(int i=0; i<N_GOAL_LINES; i++){
        if(MH_0.getMappingAt(i) == MH_1.getMappingAt(i)){
            mappings_merge[i] = MH_0.getMappingAt(i);
        } else {
            mappings_merge[i] = max(MH_0.getMappingAt(i), MH_1.getMappingAt(i));
        }
        evidence_merge[i] = MH_0.getEvidenceAt(i) + MH_1.getEvidenceAt(i);
        if(mappings_merge[i] != -1)
            n_lines_merge++;
    }
    int n_lines_in_0 = MH_0.getN_lines();
    int n_lines_in_1 = MH_1.getN_lines();
    int next_ind = n_lines_in_0;
    for(int j=0; j<n_lines_in_1; j++){
        Line goal_line = MH_1.getG_lineAt(j);
        int goal_ind = goal_line.getId();
        if(MH_0.getMappingAt(goal_ind)==-1){
            g_lines_merge[next_ind] = MH_1.getG_lineAt(j);
            c_lines_merge[next_ind] = MH_1.getC_lineAt(j);
            next_ind++;
        }
    }
    MH MH_merged(evidence_merge, mappings_merge, g_lines_merge, c_lines_merge, n_lines_merge);
    return MH_merged;
}

/**
*  Defines the comparison of two MHs based off of their scores.
*/
bool compareMHs(MH MH1, MH MH2){
    return MH1.getScore() > MH2.getScore();
}

/**
*  Sets up the line information relative to the attribute means from their image.
*/
double *calculateRels(Line l, double *a_means){
    double *rels = new double[4];
    int group_mid = (int)a_means[0];
    int l_mid = l.getPoints()[2];
    int *points_group = addrToCoords(group_mid);
    int *points_l = addrToCoords(l_mid);
    rels[0] = l.getLength()/max(pointDist(l_mid, group_mid),0.001);
    rels[1] = abs(l.getOrientation()-a_means[1]);
    rels[2] = l.getLength()/a_means[2];
    rels[3] = atan2((double)(points_l[1]-points_group[1]),(double)(points_l[0]-points_group[0]));
    return rels;
}

/**
*  Compares the relative orientational info for a line from the goal and a line from
*  the current image.
*/
double relOrientComp(double orientation1, double orientation2){
    double d_orient = abs(orientation1 - orientation2);
    if(d_orient>(PI/2)){
        orientation1 += ((orientation1 < (PI/2)) ? PI : -1*PI);
        d_orient = abs(orientation1 - orientation2);
    }
    double score = 1 - (d_orient/(PI/2));
    return score;
}

/**
*  Compares the relative distance info for a line from the goal and a line from
*  the current image. Also using this for length comparisons.
*/
double relDistComp(double dist1, double dist2){
    double d_dist_max = MAX_LENGTH/MIN_LENGTH - MIN_LENGTH/MAX_LENGTH;
    double d_dist = abs(dist1-dist2);
    double score = 1 - (d_dist/d_dist_max);
    return score;
}

/**
*  Compares the positional info for a line from the goal and a line from
*  the current image. The position refers to the angle between a line drawn 
*  from the line midpoint to the mean position for the the image's lines and 
*  the horizontal (-PI through PI). If that makes sense...
*/
double relPosComp(double pos1, double pos2){
    pos1 += PI;
    pos2 += PI;
    double d_orient = abs(pos1-pos2);
    if(d_orient>PI){
        pos1 += ((pos1 < PI) ? 2*PI : -2*PI);
        d_orient = abs(pos1-pos2);
    }
    double score = 1 - (d_orient/PI);
    return score;
}

/**
*  Brings together all of the scoring functions above and takes the average.
*/
double lineSimScore(Line l_g, Line l_c, double *a_means_g, double *a_means_c){  
    double *rels_g = calculateRels(l_g, a_means_g);
    double *rels_c = calculateRels(l_c, a_means_c);
    double score_dist = relDistComp(rels_g[0], rels_c[0]);
    double score_orient = relOrientComp(rels_g[1], rels_c[1]);
    double score_length = relDistComp(rels_g[2], rels_c[2]);
    double score_position = relPosComp(rels_g[3], rels_c[3]);
    double score = (0.2*score_dist + 0.5*score_orient + 0.25*score_length + 0.05*score_position);
    //cout << score << endl;
    return score;
}

/**
*  Populates the given MH array H by looking at the relationship matrices in the two
*  LineGroups and mapping all possible combinations within groups (eg. adj -> adj).
*  Note that this could be improved by using feedforward information from previous 
*  mappings. I'll look into this.
*/
int findMHs(MH *H, LineGroup L_g, LineGroup L_c){
    int n_MHs = 0; 
    //adjacencies
    bool *R_g = L_g.getR_left();    
    bool *R_c = L_c.getR_left();
    int belief = 0;  
    char case_name[8];
    for(int k=0; k<5; k++){
        switch(k){
            case 0:
                belief = B_LEFT;
                sprintf(case_name, "left");
                break;
            case 1:
                belief = B_ABOVE;
                R_g = L_g.getR_above();
                R_c = L_c.getR_above();
                sprintf(case_name, "above");
                break;
            case 2:
                belief = B_PAR;
                R_g = L_g.getR_par();
                R_c = L_c.getR_par();
                sprintf(case_name, "par");
                break;
            case 3:
                belief = B_PERP;
                R_g = L_g.getR_perp();
                R_c = L_c.getR_perp();
                sprintf(case_name, "perp");
                break;
            case 4:
                belief = B_ADJ;
                R_g = L_g.getR_adj();
                R_c = L_c.getR_adj();
                sprintf(case_name, "adj");
                break;
        }
        int n_r = 0;
        for(int i=0; i<(N_GOAL_LINES*N_GOAL_LINES); i++){
            if(R_g[i]){
                for(int j=0; j<(N_CURR_LINES*N_CURR_LINES); j++){
                    if(R_c[j]){
                        int *g_inds = L_g.relLines(i);
                        int *c_inds = L_c.relLines(j);
                        MH newMH(belief, L_g.getLineAt(g_inds[0]), L_g.getLineAt(g_inds[1]), L_c.getLineAt(c_inds[0]), L_c.getLineAt(c_inds[1]));
                        if(newMH.getScore() > SCORE_THRESH){
                            H[n_MHs] = newMH;
                            n_MHs++;
                            n_r++;
                        }
                        if(k>1){ //reciprocal
                            MH newMH(belief, L_g.getLineAt(g_inds[1]), L_g.getLineAt(g_inds[0]), L_c.getLineAt(c_inds[0]), L_c.getLineAt(c_inds[1]));
                            if(newMH.getScore() > SCORE_THRESH){
                                H[n_MHs] = newMH;
                                n_MHs++;
                                n_r++;
                            }
                        }
                    }
                }
            }
        }
        if(TESTING) cout << "Number in case '" << case_name << "':\t" << n_r << endl;
    }
    return n_MHs;
}

/**
*  This step uses the 'GreedyMerge' algorithm outlined in Forbus and Oblinger, 1990.
*  This allows us to come up with an approximation for the best full mappings, without
*  running into the problem of combinatorial explosion. To see why this would happen,
*  refer to the discussion of the merge step in Falkenhainer, Forbus and Gentner, 1989,
*  and see their explanation of the worst-case type of system.
*/
int mergeMHs(MH *H, int n_MHs, GMap *G, LineGroup L_g, LineGroup L_c){
    int n_checked = 0;
    vector<MH> H_full;
    vector<MH> H_temp;
    for(int i=0; i<n_MHs; i++){
        H_temp.push_back(H[i]);
    }
    sort(H_temp.begin(), H_temp.end(), compareMHs);  //big to small
    bool *is_used = new bool[H_temp.size()];
    fill_n(is_used,H_temp.size(),0);
    int n_full = 0;
    for(int i=0; i<NUM_TO_MAP; i++){
        bool out_of_seeds = 1;
        MH MH_snowball;
        int seed_ind = -1;
        for(int j=0; j<H_temp.size(); j++){
            if(!is_used[j]){
                MH_snowball = H_temp.at(j);
                seed_ind = j;
                is_used[j] = 1;
                out_of_seeds = 0;
                break;
            }
        }
        if(out_of_seeds) break;
        for(int k=0; k<H_temp.size(); k++){
            if(k==seed_ind) continue;
            if(canMerge(MH_snowball, H_temp.at(k))){
                MH MH_merged = mergeMHs(MH_snowball, H_temp.at(k));
                MH_snowball = MH_merged;
            }
        }        
        H_full.push_back(MH_snowball);
        n_full++;
    }
    sort(H_full.begin(), H_full.end(), compareMHs);  //big to small
    for(int i=0; i<n_full; i++){
        GMap newGMap(H_full.at(i).getN_lines(), H_full.at(i).getScore(), H_full.at(i).getMappings());
        G[i] = newGMap;
    }
    return n_full;
}

//--------------------GMaps-------------------//

/**
*  This type is very similar to MH. Is it necessary?
*/
GMap::GMap(int n_lines, double MHScore, int *mappings){
    this->n_lines = n_lines;
    score = MHScore;
    order_g = new int[n_lines];
    order_c = new int[n_lines];
    int ind_next = 0;
    for(int i=0; i<N_GOAL_LINES; i++){
        if(mappings[i] != -1){
            order_g[ind_next] = i;
            order_c[ind_next] = mappings[i];
            ind_next++;
        }
    }
}

int GMap::getN_lines(){
    return n_lines;
}

double GMap::getScore(){
    return score;
}

int *GMap::getOrder_g(){
    return order_g;
}

int *GMap::getOrder_c(){
    return order_c;
}

/**
*  Given the best GMaps and some other info on the lines etc., loops over and draws corresponding lines.
*  Lines mapped to each other have the same color between images. There are only 12 color options right now,
*  so if more than 12 lines are mapped, it could get messy.
*/
void drawGMaps(char *log_path, int numMaps, GMap *mappings, LineGroup L_g, LineGroup L_c, Mat cleanGoal, Mat cleanCurr){
    char map_name[100];
    Mat goalDraw;
    Mat currDraw;
    Mat mapOut(YDIM, 2*XDIM, cleanGoal.type());
    Mat ROI;
    for(int i=0; i<numMaps; i++){
        sprintf(map_name,"%smapping%d.jpg",log_path,i);
	GMap map = mappings[i];
        if(TESTING) cout << "Mapping: " << i << "\t (score = " << map.getScore() << ")" << endl;
	goalDraw = cleanGoal.clone(); 
	currDraw = cleanCurr.clone();
	int *gKey = map.getOrder_g();
	int *cKey = map.getOrder_c();
	int n_lines = map.getN_lines();
	for(int j=0; j<n_lines; j++){
		int gInd = gKey[j];
		int cInd = cKey[j];
                if(TESTING) cout << "Mapped are: " << gInd << " and " << cInd << endl;
		Line goalLine = L_g.getLineAt(gInd);
		Line currLine = L_c.getLineAt(cInd);
		int *goalPoints = goalLine.getPoints();
		int *goalEnd1 = addrToCoords(goalPoints[0]);
		int *goalEnd2 = addrToCoords(goalPoints[1]);
		Point goalpt1 = Point(goalEnd1[0],goalEnd1[1]);
		Point goalpt2 = Point(goalEnd2[0],goalEnd2[1]);
		line(goalDraw, goalpt1, goalpt2, lineColors[j%12], 3);
		int *currPoints = currLine.getPoints();
		int *currEnd1 = addrToCoords(currPoints[0]);
		int *currEnd2 = addrToCoords(currPoints[1]);
		Point currpt1 = Point(currEnd1[0],currEnd1[1]);
		Point currpt2 = Point(currEnd2[0],currEnd2[1]);
		line(currDraw, currpt1, currpt2, lineColors[j%12], 3);
	}
	ROI = mapOut(Rect(0,0,XDIM,YDIM));
	goalDraw.copyTo(ROI);
	ROI = mapOut(Rect(XDIM,0,XDIM,YDIM));
	currDraw.copyTo(ROI);
        Point div1 = Point(XDIM, 0);
        Point div2 = Point(XDIM, YDIM);
        line(mapOut, div1, div2, lineColors[0], 2);
        imwrite(map_name, mapOut);
    }  
}


#if 1
/**
*  Returns the orientation to which lines should be compared based on
*  the longest line in the goal.
*/
double getValByLongLine(int n_lines, int *goalOrder, int *currOrder, LineGroup L_g, LineGroup L_c){
    double longest = 0.0;
    double retVal = 0.0;
    for(int i=0; i<n_lines; i++){
        Line gLineAtI = L_g.getLineAt(goalOrder[i]);
        if(gLineAtI.getLength()>longest){
            longest = gLineAtI.getLength();
            Line cLineAtI = L_c.getLineAt(currOrder[i]);
            double orientation1 = gLineAtI.getOrientation();
            double orientation2 = cLineAtI.getOrientation();
            retVal = abs(orientation1 - orientation2);
            if(retVal>(PI/2)){
                orientation1 += ((orientation1 < (PI/2)) ? PI : -1*PI);
                retVal = abs(orientation1 - orientation2);
            }
        }
    }
    return retVal;
}

/**
*  Returns the orientation to which lines should be compared based on
*  the best scoring line match.
*/
double getValByBestMatch(int n_lines, int *goalOrder, int *currOrder, LineGroup L_g, LineGroup L_c){
    double bestScore = 0.0;
    double retVal = 0.0;
    for(int i=0; i<n_lines; i++){
        Line gLineAtI = L_g.getLineAt(goalOrder[i]);
        Line cLineAtI = L_c.getLineAt(currOrder[i]);
        double score = lineSimScore(gLineAtI, cLineAtI, G_MEANS, C_MEANS);
        if(score>bestScore){
            bestScore = score;
            double orientation1 = gLineAtI.getOrientation();
            double orientation2 = cLineAtI.getOrientation();
            retVal = abs(orientation1 - orientation2);
            if(retVal>(PI/2)){
                orientation1 += ((orientation1 < (PI/2)) ? PI : -1*PI);
                retVal = abs(orientation1 - orientation2);
            }
        }
    }
    return retVal;
}

/**
*  Returns the orientation to which lines should be compared based on
*  the average orientation in the goal.
*/
double getAvgVal(int n_lines, int *order_g, int *order_c, LineGroup L_g, LineGroup L_c){
    double totOrient = 0.0;
    for(int i=0; i<n_lines; i++){
        Line gLineAtI = L_g.getLineAt(order_g[i]);
        Line cLineAtI = L_c.getLineAt(order_c[i]);
        double orientation1 = gLineAtI.getOrientation();
        double orientation2 = cLineAtI.getOrientation();
        double orient = abs(orientation1 - orientation2);
        if(orient>(PI/2)){
            orientation1 += ((orientation1 < (PI/2)) ? PI : -1*PI);
            orient = abs(orientation1 - orientation2);
        }
        totOrient += orient;
    }
    double retVal = totOrient/n_lines;
    return retVal;
}

/**
*  In order to ensure we have a good match, we want to verify the lines. There are three methods-
*      0) Compare to longest line
*      1) Compare to best scoring line match
*      2) Compare to avg offset
*  Returns the index of the top match based on score and verification.
*/
int getVerifiedBestMap(int numMaps, GMap *mappings, LineGroup L_g, LineGroup L_c, int method){
    for(int i=0; i<numMaps; i++){
        GMap thisGMap = mappings[i];
        int numLinesMapped = thisGMap.getN_lines();
        int *Lines_g = thisGMap.getOrder_g();
        int *Lines_c = thisGMap.getOrder_c();
        double compareVal = 0.0;
        switch(method){
            case 0:
                compareVal = getValByLongLine(numLinesMapped, Lines_g, Lines_c, L_g, L_c);
                break;
            case 1:
                compareVal = getValByBestMatch(numLinesMapped, Lines_g, Lines_c, L_g, L_c);
                break;
            case 2:
                compareVal = getAvgVal(numLinesMapped, Lines_g, Lines_c, L_g, L_c);
                break;
        }
        int numMissed = 0;
        for(int j=0; j<numLinesMapped; j++){
            Line goalLine = L_g.getLineAt(Lines_g[j]);
            Line currLine = L_c.getLineAt(Lines_c[j]);
            double orientation1 = goalLine.getOrientation();
            double orientation2 = currLine.getOrientation();
            double offset = abs(orientation1 - orientation2);
            if(offset>(PI/2)){
                orientation1 += ((orientation1 < (PI/2)) ? PI : -1*PI);
                offset = abs(orientation1 - orientation2);
            }
            if(abs(offset-compareVal)>VERIF_THRESH){
                numMissed++;
                cout << "The offset is: " << offset << ", and the compareVal is: " << compareVal << endl;
            }
        }
        cout << "Map #: " << i << "; NumMissed = " << numMissed << endl;
    }
    return 0;
}
#endif



